package ai.kastrax.examples.agent

import ai.kastrax.core.agent.AgentStreamOptions
import ai.kastrax.core.agent.agent
import ai.kastrax.core.agent.architecture.ReflectiveAgent
import ai.kastrax.core.agent.architecture.reflectiveAgent
import ai.kastrax.integrations.deepseek.deepSeek
import ai.kastrax.integrations.deepseek.DeepSeekModel
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.flow.collect

/**
 * 反思型Agent示例
 */
fun main() = runBlocking {
    // 创建基础LLM
    val llm = deepSeek {
        // 直接设置 API 密钥
        apiKey("sk-85e83081df28490b9ae63188f0cb4f79")

        // 设置模型
        model(DeepSeekModel.DEEPSEEK_CHAT)

        // 设置生成参数
        temperature(0.7)
        maxTokens(2000)
        topP(0.95)
    }

    // 使用DSL创建基础Agent
    val baseAgent = agent {
        name = "基础Agent"
        model = llm
    }

    // 使用DSL创建反思型Agent
    val reflectiveAgent = reflectiveAgent {
        baseAgent(baseAgent)
        config {
            enablePreReflection(true)
            enableResponseReflection(true)
            enablePostReflection(true)
            enableLearningFromReflection(true)
        }
    }

    // 注意：在新的实现中，没有addLesson方法
    // 反思型Agent会自动从交互中学习

    println("=== 反思型Agent示例 ===")

    // 创建会话
    val sessionId = "session-123"

    // 第一个问题
    val prompt1 = "什么是大语言模型？它们是如何工作的？"
    println("\n问题1: $prompt1")

    val options1 = AgentStreamOptions()
    val optionsWithMetadata1 = options1.copy(metadata = mapOf("sessionId" to sessionId))
    
    try {
        val response1 = reflectiveAgent.stream(
            prompt = prompt1,
            options = optionsWithMetadata1
        )

        println("回答1:")
        
        // 处理流式文本响应
        response1.textStream?.collect { chunk ->
            print(chunk)
            kotlinx.coroutines.delay(25)
        } ?: run {
            print(response1.text)
        }
        
        println() // 换行
        
    } catch (e: Exception) {
        println("生成响应时发生错误: ${e.message}")
    }

    // 第二个问题
    val prompt2 = "大语言模型面临哪些挑战和局限性？"
    println("\n问题2: $prompt2")

    val options2 = AgentStreamOptions()
    val optionsWithMetadata2 = options2.copy(metadata = mapOf("sessionId" to sessionId))
    
    try {
        val response2 = reflectiveAgent.stream(
            prompt = prompt2,
            options = optionsWithMetadata2
        )

        println("回答2:")
        
        // 处理流式文本响应
        response2.textStream?.collect { chunk ->
            print(chunk)
            kotlinx.coroutines.delay(25)
        } ?: run {
            print(response2.text)
        }
        
        println() // 换行
        
    } catch (e: Exception) {
        println("生成响应时发生错误: ${e.message}")
    }

    // 查看会话的反思记录
    val reflections = reflectiveAgent.getSessionReflections(sessionId)
    println("\n=== 会话反思记录 ===")
    println("共有 ${reflections.size} 条反思记录")

    // 显示部分反思记录
    if (reflections.isNotEmpty()) {
        val sample = reflections.take(2)
        sample.forEachIndexed { index, reflection ->
            println("\n反思记录 ${index + 1}:")
            println("类型: ${reflection.type}")
            println("内容: ${reflection.content}")
        }
    }

    // 查看全局反思
    val globalReflections = reflectiveAgent.getGlobalReflections()
    println("\n=== 全局反思 ===")
    println("共有 ${globalReflections.size} 条全局反思")
}
